﻿using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Diagnostics.Contracts;
using Rxx;

namespace System.Net.Sockets
{
	public static partial class SocketExtensions
	{
		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginConnect(this Socket socket, EndPoint remoteEP, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				RemoteEndPoint = remoteEP
			};

			return BeginInvoke(callback, state, args, socket.ConnectAsync);
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginConnect(this Socket socket, IPAddress address, int port, AsyncCallback requestCallback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				RemoteEndPoint = new IPEndPoint(address, port)
			};

			return BeginInvoke(requestCallback, state, args, socket.ConnectAsync);
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginConnect(this Socket socket, string host, int port, AsyncCallback requestCallback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				RemoteEndPoint = new DnsEndPoint(host, port)
			};

			return BeginInvoke(requestCallback, state, args, socket.ConnectAsync);
		}

		internal static void EndConnect(this Socket socket, IAsyncResult asyncResult)
		{
			Contract.Requires(asyncResult != null);

			using (var args = EndInvoke(asyncResult))
			{
				// do nothing
			}
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginReceive(this Socket socket, IList<ArraySegment<byte>> buffers, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				BufferList = buffers
			};

			return BeginInvoke(callback, state, args, socket.ReceiveAsync);
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginReceive(this Socket socket, IList<ArraySegment<byte>> buffers, out SocketError errorCode, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				BufferList = buffers
			};

			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);

			if (result.CompletedSynchronously)
			{
				errorCode = args.SocketError;

				// Prevent the error from being thrown in EndInvoke
				result.OutOfBandData = new object();
			}
			else
			{
				errorCode = default(SocketError);
			}

			return result;
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginReceive(this Socket socket, byte[] buffer, int offset, int size, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs();

			args.SetBuffer(buffer, offset, size);

			return BeginInvoke(callback, state, args, socket.ReceiveAsync);
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginReceive(this Socket socket, byte[] buffer, int offset, int size, out SocketError errorCode, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs();

			args.SetBuffer(buffer, offset, size);

			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);

			if (result.CompletedSynchronously)
			{
				errorCode = args.SocketError;

				// Prevent the error from being thrown in EndInvoke
				result.OutOfBandData = new object();
			}
			else
			{
				errorCode = default(SocketError);
			}

			return result;
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginReceiveFrom(this Socket socket, byte[] buffer, int offset, int size, ref EndPoint remoteEP, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				RemoteEndPoint = remoteEP
			};

			args.SetBuffer(buffer, offset, size);

			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);

			if (result.CompletedSynchronously)
			{
				remoteEP = args.RemoteEndPoint;
			}

			return result;
		}

		internal static int EndReceive(this Socket socket, IAsyncResult asyncResult)
		{
			Contract.Requires(asyncResult != null);

			using (var args = EndInvoke(asyncResult))
			{
				return args.BytesTransferred;
			}
		}

		[SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket",
			Justification = "Needed for parity with .NET 4.0 Framework.")]
		internal static int EndReceive(this Socket socket, IAsyncResult asyncResult, out SocketError errorCode)
		{
			Contract.Requires(asyncResult != null);

			using (var args = EndInvoke(asyncResult))
			{
				errorCode = args.SocketError;

				return args.BytesTransferred;
			}
		}

		[SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket",
			Justification = "Needed for parity with .NET 4.0 Framework.")]
		internal static int EndReceiveFrom(this Socket socket, IAsyncResult asyncResult, ref EndPoint endPoint)
		{
			Contract.Requires(asyncResult != null);

			using (var args = EndInvoke(asyncResult))
			{
				endPoint = args.RemoteEndPoint;

				return args.BytesTransferred;
			}
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginSend(this Socket socket, IList<ArraySegment<byte>> buffers, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				BufferList = buffers
			};

			return BeginInvoke(callback, state, args, socket.SendAsync);
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginSend(this Socket socket, IList<ArraySegment<byte>> buffers, out SocketError errorCode, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				BufferList = buffers
			};

			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);

			if (result.CompletedSynchronously)
			{
				errorCode = args.SocketError;

				// Prevent the error from being thrown in EndInvoke
				result.OutOfBandData = new object();
			}
			else
			{
				errorCode = default(SocketError);
			}

			return result;
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginSend(this Socket socket, byte[] buffer, int offset, int size, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs();

			args.SetBuffer(buffer, offset, size);

			return BeginInvoke(callback, state, args, socket.SendAsync);
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginSend(this Socket socket, byte[] buffer, int offset, int size, out SocketError errorCode, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs();

			args.SetBuffer(buffer, offset, size);

			var result = BeginInvoke(callback, state, args, socket.ReceiveAsync);

			if (result.CompletedSynchronously)
			{
				errorCode = args.SocketError;

				// Prevent the error from being thrown in EndInvoke
				result.OutOfBandData = new object();
			}
			else
			{
				errorCode = default(SocketError);
			}

			return result;
		}

		[SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope",
			Justification = "The args object is disposed in the EndInvoke method.")]
		internal static IAsyncResult BeginSendTo(this Socket socket, byte[] buffer, int offset, int size, EndPoint remoteEP, AsyncCallback callback, object state)
		{
			Contract.Requires(socket != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var args = new SocketAsyncEventArgs()
			{
				RemoteEndPoint = remoteEP
			};

			args.SetBuffer(buffer, offset, size);

			return BeginInvoke(callback, state, args, socket.SendAsync);
		}

		internal static int EndSend(this Socket socket, IAsyncResult asyncResult)
		{
			Contract.Requires(asyncResult != null);

			using (var args = EndInvoke(asyncResult))
			{
				return args.BytesTransferred;
			}
		}

		[SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket",
			Justification = "Needed for parity with .NET 4.0 Framework.")]
		internal static int EndSend(this Socket socket, IAsyncResult asyncResult, out SocketError errorCode)
		{
			Contract.Requires(asyncResult != null);

			using (var args = EndInvoke(asyncResult))
			{
				errorCode = args.SocketError;

				return args.BytesTransferred;
			}
		}

		[System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", MessageId = "socket", 
			Justification = "Must have the same signature as in .NET 4.0.")]
		internal static int EndSendTo(this Socket socket, IAsyncResult asyncResult)
		{
			Contract.Requires(asyncResult != null);

			using (var args = EndInvoke(asyncResult))
			{
				return args.BytesTransferred;
			}
		}

		private static AsyncResult<SocketAsyncEventArgs> BeginInvoke(AsyncCallback callback, object state, SocketAsyncEventArgs args, Func<SocketAsyncEventArgs, bool> target)
		{
			Contract.Requires(args != null);
			Contract.Requires(target != null);
			Contract.Ensures(Contract.Result<IAsyncResult>() != null);

			var result = new AsyncResult<SocketAsyncEventArgs>(state);

			EventHandler<SocketAsyncEventArgs> completed = (sender, e) =>
			{
				if (!result.IsCompleted)
				{
					result.Complete(args, synchronous: false);
				}

				if (callback != null)
				{
					callback(result);
				}
			};

			args.Completed += completed;

			if (!target(args))
			{
				result.Complete(args, synchronous: true);

				completed(null, null);
			}

			return result;
		}

		private static SocketAsyncEventArgs EndInvoke(IAsyncResult asyncResult)
		{
			Contract.Requires(asyncResult != null);
			Contract.Ensures(Contract.Result<SocketAsyncEventArgs>() != null);

			var result = (AsyncResult<SocketAsyncEventArgs>) asyncResult;

			using (result)
			{
				result.Wait();

				var args = result.Data;

				Contract.Assume(args != null);

				if (result.OutOfBandData == null)
				{
					if (args.ConnectByNameError != null)
					{
						throw args.ConnectByNameError;
					}

					if (args.SocketError != SocketError.Success)
					{
						throw new SocketException((int) args.SocketError);
					}
				}

				return args;
			}
		}
	}
}